import cv2
import numpy as np

def extract_features(image):
    """提取特征点并计算描述符"""
    img = cv2.imread(image)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    sift = cv2.SIFT_create()
    keypoints, descriptors = sift.detectAndCompute(gray, None)
    return keypoints, descriptors

def match_method():
    return cv2.BFMatcher(crossCheck=True)

def match_features(desc1, desc2):
    """利用近邻的特征点匹配方法进行特征点匹配"""
    bf = cv2.BFMatcher(cv2.NORM_L2, crossCheck=True)
    matches = bf.match(desc1, desc2)
    matches = sorted(matches, key=lambda x: x.distance)
    return matches

def filter_matches_by_fundamental_matrix(kp1, kp2, matches):
    """利用几何一致性过滤误匹配，估计基础矩阵F"""
    pts1 = np.float32([kp1[m.queryIdx].pt for m in matches])
    pts2 = np.float32([kp2[m.trainIdx].pt for m in matches])
    F, mask = cv2.findFundamentalMat(pts1, pts2, cv2.FM_RANSAC)
    matches = [m for m, inlier in zip(matches, mask) if inlier]
    return matches, F , pts1, pts2

def filter_matches_by_homography(kp1, kp2, matches):
    """利用几何一致性过滤误匹配，估计单应矩阵H"""
    pts1 = np.float32([kp1[m.queryIdx].pt for m in matches])
    pts2 = np.float32([kp2[m.trainIdx].pt for m in matches])
    H, mask = cv2.findHomography(pts1, pts2, cv2.RANSAC)
    matches = [m for m, inlier in zip(matches, mask) if inlier]
    return matches,H,pts1,pts2

def preprocess_images(image1, image2):
    """主函数：处理图像集并输出几何效验后的特征点匹配结果"""
    kp1, desc1 = extract_features(image1)
    kp2, desc2 = extract_features(image2)
    
    matches = match_features(desc1, desc2)
    # 基本矩阵
    matches_f, F, pts1_f, pts2_f = filter_matches_by_fundamental_matrix(kp1, kp2, matches)
    # 单应矩阵
    matches_h, H, pts1_h, pts2_h = filter_matches_by_homography(kp1, kp2, matches)
    
    if len(matches_f) > len(matches_h):
        return matches_f,F,pts1_f,pts2_f
    else:
        return matches_h,H,pts1_h,pts2_h

if __name__=='__main__':
    # 示例用法
    image1 = cv2.imread(r'data\castle-P19\images\0000.jpg')
    image2 = cv2.imread(r'data\castle-P19\images\0001.jpg')
    matches,Matrix, = preprocess_images(image1, image2)
